##########################################################################################
# Title: Figure 2 — Global Gradients (4 Panels): Pop Totals, Density, Share, Cumulative
# Section: 2 (Figures)
# Purpose: Aggregate 100m GHSL grid cells within 20 km of BDs to plot global benchmarks
# IO:
#   IN : Processed_Data/Intermediate/Section_2/Gradients/gams_gradient.rds
#   OUT: Graphs/Figure_2.png
##########################################################################################

rm(list = ls()); options(stringsAsFactors = FALSE)

suppressPackageStartupMessages({
  library(ggplot2)
  library(dplyr)
  library(patchwork)
  library(grid)   # unit()
})

# ---- Project root (align with run_all.R) -------------------------------------
if (!requireNamespace("here", quietly = TRUE)) install.packages("here", quiet = TRUE)
ROOT      <- normalizePath(here::here(), winslash = "/")
GRAD_DIR  <- file.path(ROOT, "Processed_Data", "Intermediate", "Section_2", "Gradients")
GRAPH_DIR <- file.path(ROOT, "Graphs")
dir.create(GRAPH_DIR, showWarnings = FALSE, recursive = TRUE)

# ---- Load --------------------------------------------------------------------
final_rds <- file.path(GRAD_DIR, "gams_gradient.rds")
if (!file.exists(final_rds)) stop("Missing file: ", final_rds)
df <- readRDS(final_rds)

req_cols <- c("city_name","GHS_POP_100_2020","min_dist_cbd","south_africa","other_africa")
if (!all(req_cols %in% names(df))) stop("Required columns missing in gams_gradient.rds")

# ---- Transform ---------------------------------------------------------------
df <- df %>%
  mutate(
    dist_km = min_dist_cbd / 1000,
    # handle both logical and 0/1 with NA-safety
    south   = dplyr::coalesce(as.logical(south_africa), FALSE),
    other   = dplyr::coalesce(as.logical(other_africa), FALSE),
    group   = dplyr::case_when(
      south ~ "South Africa",
      other ~ "Other African Cities",
      TRUE  ~ "Cities Outside Africa"
    ),
    # GHSL 100m pixel density → per km² scale factor = 100
    pop_per_km2 = GHS_POP_100_2020 * 100
  )

# ---- Bin distances (1 km bands up to 20 km) ---------------------------------
bin_width <- 1
df <- df %>%
  mutate(
    dist_bin = cut(dist_km, breaks = seq(0, 20, by = bin_width),
                   include.lowest = TRUE, right = FALSE)
  ) %>%
  filter(!is.na(dist_bin))

# ---- City counts per group ---------------------------------------------------
city_counts <- df %>%
  dplyr::distinct(group, city_name) %>%
  dplyr::count(group, name = "n_cities")

# ---- Aggregate to bands ------------------------------------------------------
agg_df <- df %>%
  group_by(group, dist_bin) %>%
  summarise(
    avg_density   = mean(pop_per_km2, na.rm = TRUE),
    total_pop_raw = sum(GHS_POP_100_2020, na.rm = TRUE),
    .groups = "drop"
  ) %>%
  left_join(city_counts, by = "group") %>%
  mutate(
    total_pop = total_pop_raw / n_cities
  ) %>%
  group_by(group) %>%
  mutate(
    cumulative_share = cumsum(total_pop) / sum(total_pop),
    normalized_share = total_pop / sum(total_pop)
  ) %>%
  ungroup()

# ---- Order groups ------------------------------------------------------------
agg_df$group <- factor(
  agg_df$group,
  levels = c("Cities Outside Africa", "Other African Cities", "South Africa")
)

# ---- Bin midpoints (for plotting) -------------------------------------------
agg_df <- agg_df %>%
  mutate(dist_mid = as.numeric(gsub("\\[|,.*", "", dist_bin)) + 0.5)

# ---- Aesthetics --------------------------------------------------------------
plot_colors <- c("Cities Outside Africa" = "grey50",
                 "Other African Cities"  = "black",
                 "South Africa"          = "red")
plot_lw    <- c("Cities Outside Africa" = 3,
                "Other African Cities"  = 1,
                "South Africa"          = 1)
plot_alpha <- c("Cities Outside Africa" = 0.4,
                "Other African Cities"  = 1,
                "South Africa"          = 1)
base_font_size <- 12

legend_theme <- theme(
  legend.position   = "bottom",
  legend.title      = element_blank(),
  legend.text       = element_text(size = 14),
  legend.key.width  = unit(2.5, "lines"),
  legend.key.height = unit(1.5, "lines"),
  legend.spacing.x  = unit(1.2, "cm")
)

legend_guides <- guides(
  color = guide_legend(
    title = NULL,
    override.aes = list(
      linewidth = as.numeric(plot_lw)[match(levels(agg_df$group), names(plot_lw))],
      alpha     = as.numeric(plot_alpha)[match(levels(agg_df$group), names(plot_alpha))]
    )
  ),
  linewidth = "none",
  alpha     = "none"
)

# ---- Panels ------------------------------------------------------------------
p1 <- ggplot(agg_df, aes(x = dist_mid, y = total_pop / 1e6, group = group, color = group)) +
  geom_line(aes(linewidth = group, alpha = group)) +
  scale_color_manual(values = plot_colors) +
  scale_linewidth_manual(values = plot_lw) +
  scale_alpha_manual(values = plot_alpha) +
  scale_x_continuous(breaks = 1:20) +
  labs(title = "A. Average Total Population per City",
       x = "Distance to Business District (km)",
       y = "Population (millions)") +
  theme_minimal(base_size = base_font_size) +
  theme(panel.grid = element_blank(), axis.line = element_line(), legend.position = "none")

p2 <- ggplot(agg_df, aes(x = dist_mid, y = avg_density / 1000, group = group, color = group)) +
  geom_line(aes(linewidth = group, alpha = group)) +
  scale_color_manual(values = plot_colors) +
  scale_linewidth_manual(values = plot_lw) +
  scale_alpha_manual(values = plot_alpha) +
  scale_x_continuous(breaks = 1:20) +
  labs(title = "B. Average Population Density",
       x = "Distance to Business District (km)",
       y = "Avg. Density (000/km²)") +
  theme_minimal(base_size = base_font_size) +
  theme(panel.grid = element_blank(), axis.line = element_line(), legend.position = "none")

p3 <- ggplot(agg_df, aes(x = dist_mid, y = normalized_share, group = group, color = group)) +
  geom_line(aes(linewidth = group, alpha = group)) +
  scale_color_manual(values = plot_colors) +
  scale_linewidth_manual(values = plot_lw) +
  scale_alpha_manual(values = plot_alpha) +
  scale_x_continuous(breaks = 1:20) +
  labs(title = "C. Share of Population",
       x = "Distance to Business District (km)",
       y = "Share of Population") +
  theme_minimal(base_size = base_font_size) +
  theme(panel.grid = element_blank(), axis.line = element_line(), legend.position = "none")

p4 <- ggplot(agg_df, aes(x = dist_mid, y = cumulative_share, group = group, color = group)) +
  geom_line(aes(linewidth = group, alpha = group)) +
  scale_color_manual(values = plot_colors) +
  scale_linewidth_manual(values = plot_lw) +
  scale_alpha_manual(values = plot_alpha) +
  scale_x_continuous(breaks = 1:20) +
  labs(title = "D. Cumulative Population Share",
       x = "Distance to Business District (km)",
       y = "Cumulative Population Share") +
  theme_minimal(base_size = base_font_size) +
  theme(panel.grid = element_blank(), axis.line = element_line())

# ---- Combine + Save ----------------------------------------------------------
final_plot <- (p1 + p2) / (p3 + p4) +
  plot_layout(guides = "collect") &
  legend_theme &
  legend_guides

out_file <- file.path(GRAPH_DIR, "Figure_2.png")
ggsave(out_file, final_plot, width = 7, height = 8, dpi = 300)
message("✅ Saved: ", normalizePath(out_file))

